from __future__ import annotations
from typing import Callable, TYPE_CHECKING, Union, Any, TypeVar, Type
from shared_enum import Layer, Process
from hashtool import Hash
from uuid import uuid4

T = TypeVar('T')

if TYPE_CHECKING:
    from frame_manager import FrameManager
    from neo4j.graph import Node, Relationship
    from neo_manager import NeoManager

'''
frame processing
0 uncreated:  no marker, no uuid, no hash
1 filled:  has marker, no uuid, no hash
2 checked:  has marker(never use marker after stage checked! use uuid instead), has uuid, hash hash, pass check
3 stored:  optional has marker(may same as uuid), has correct uuid, has hash
'''

class Frame:
    def __init__(self, frame_manager: FrameManager, frame_name: str) -> None:
        self.frame_manager = frame_manager
        self.name = frame_name
        self.script = frame_manager.frame_script_map[frame_name]
        self.layer = None
        self.elements: dict[str, ElementContainer] = None
        self.marker: str = None
        self.uuid = None
        self.hash = None
        self.hash_all = None
        self._process = Process.Uncreated
        self.corpus = None
        self.source = None

    def _hash(self) -> Frame:
        _hash = Hash(self.name)
        for element_name in sorted(self.elements):
            element_container = self.elements[element_name]
            if element_container.hash is None:
                element_container._hash()
            element_hash = element_container.hash
            _hash.update(element_name)
            _hash.update(element_hash)
        self.hash = _hash.value()
        self.hash_all = Hash(self.hash).update(self.corpus).update(self.source).value()
        return self

    def create_from_dict(self, data: dict) -> Frame:
        if data['marker'] in self.frame_manager.marker_map:
            return self.frame_manager.marker_map[data['marker']]
        else:
            self.corpus = data['corpus']
            self.source = data['source']
            self.marker = data['marker']
            self.elements = {}
            for element_name, element_schema in self.frame_manager.frame_name_map[self.name]['elements'].items():
                element = element_schema.create_from_dict(self, data)
                if element:
                    self.elements[element_name] = element
            self._process = Process.Filled
            self.layer = Layer.Annotation
            self.frame_manager.marker_map[self.marker] = self
            return self

    def create_from_db(self, neo_manager: NeoManager, frame_node: Node) -> Frame:
        uuid = frame_node['index']
        if neo_manager.layer is Layer.Annotation:
            marker_map = self.frame_manager.marker_map
        elif neo_manager.layer is Layer.Primary:
            marker_map = self.frame_manager.marker_map_primary
        else:
            marker_map = self.frame_manager.marker_map_prototype
        if uuid in marker_map:
            return marker_map[uuid]
        else:
            self.uuid = uuid
            self.marker = uuid
            self.hash = frame_node['hash']
            self.corpus = frame_node['corpus']
            self.source = frame_node['source']
            self.hash_all = Hash(self.hash).update(self.corpus).update(self.source).value()
            self.elements = {}
            for element_name, element_schema in self.frame_manager.frame_name_map[self.name]['elements'].items():
                element = element_schema.create_from_db(self, neo_manager, frame_node)
                if element:
                    self.elements[element_name] = element
            self._process = Process.Stored
            self.layer = neo_manager.layer
            marker_map[uuid] = self
            return self

    def save(self, neo_manager: NeoManager) -> None:
        layer = neo_manager.layer
        if self.layer is not layer:
            raise Exception(f'can not store {self.layer} frame in {layer} database')
        if layer is Layer.Annotation:
            marker_map = self.frame_manager.marker_map
        elif layer is Layer.Primary:
            marker_map = self.frame_manager.marker_map_primary
        elif layer is Layer.Prototype:
            marker_map = self.frame_manager.marker_map_prototype
        else:
            raise ValueError
        same_frames: list[Frame] = neo_manager.search(self.name, self.frame_manager)
        for frame in same_frames:
            if frame.hash_all == self.hash_all:
                self.uuid = frame.uuid
                marker_map[self.uuid] = self
                self._process = Process.Stored
                return None
        with neo_manager.driver.session(database=layer.name) as session:
            session.run(
                '''
                CREATE (:FrameInstance{frame: $name, index: $index, hash: $hash, corpus: $corpus, source: $source, focus: $focus})
                ''', name=self.name, index=self.uuid, hash=self.hash, corpus=self.corpus, source=self.source, focus=[]
            )
            for element_name, element in self.elements.items():
                element.save(self, neo_manager)
        self._process = Process.Stored

    def check(self, check_fun: Callable[[Frame], bool], generate_fun: Callable[[Frame], Frame]) -> None:
        if check_fun(self):
            if self.uuid is None:
                self.uuid = uuid4().hex
            if self.uuid not in self.frame_manager.marker_map:
                self.frame_manager.marker_map[self.uuid] = self
            generate_fun(self)
            self._hash()
            self._process = Process.Checked
        else:
            raise Exception(f'{self.marker}: {self.name} did not pass check')

class ElementSchema:
    def __init__(self, name: str, is_core: bool, type_restraint: list[str], container: Type[ElementContainer]):
        self.name = name
        self.is_core = is_core
        self.type_restraint = type_restraint
        self.container = container

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        return self.container(self.name, self.is_core, self.type_restraint).create_from_dict(frame, data)

    def create_from_db(self, frame: Frame, neo_manager: NeoManager, frame_node: Node) -> ElementContainer:
        return self.container(self.name, self.is_core, self.type_restraint).create_from_db(frame, neo_manager, frame_node)

    def create_empty(self) -> ElementContainer:
        return self.container(self.name, self.is_core, self.type_restraint)

class ElementContainer:
    def __init__(self, name, is_core, type_restraint) -> None:
        self.name = name
        self.is_core = is_core
        self.type_restraint = type_restraint
        self.is_focus: bool = None
        self.hash: str = None
        self.value = None
        self._process: Process = Process.Uncreated

    def _hash(self) -> ElementContainer:
        raise NotImplementedError

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        raise NotImplementedError

    def create_from_db(self, frame: Frame, neo_manager: NeoManager, frame_node: Node) -> ElementContainer:
        raise NotImplementedError

    def save(self, frame: Frame, neo_manager: NeoManager) -> None:
        raise NotImplementedError

# filter ::= Callable(is_core: bool, is_focus: bool) -> bool
    def apply(self, fun: Callable[[T], T], filter: Callable[[bool, bool], bool]=lambda i, j: True) -> ElementContainer:
        raise NotImplementedError

    @staticmethod
    def create_structure_from_db(frame_manager: FrameManager, neo_manager: NeoManager, node: Node) -> Union[Frame, Ontology]:
        if 'FrameInstance' in node.labels:
            return Frame(frame_manager, node['frame']).create_from_db(neo_manager, node)
        elif 'Ontology' in node.labels:
            return Ontology(frame_manager, node['name']).create_from_db(node)
        else:
            raise ValueError

class BaseContainer(ElementContainer):
    def __init__(self, name, is_core, type_restraint) -> None:
        super().__init__(name, is_core, type_restraint)
        self.value: Union[str, bool] = None

    def _hash(self) -> ElementContainer:
        _hash = Hash(self.name).update(self.is_focus).update(self.value)
        self.hash = _hash.value()
        return self

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        if self.name in data['element']:
            self.value = data['element'][self.name]
            self.is_focus = self.name in data['focus']
            self._process = Process.Filled
            return self
        else:
            return None

    def create_from_db(self, frame: Frame, noe_manager: NeoManager, frame_node: Node) -> ElementContainer:
        if self.name in frame_node:
            self.value = frame_node[self.name]
            self.is_focus = self.name in frame_node['focus']
            self._process = Process.Filled
            return self
        else:
            return None

    def save(self, frame: Frame, neo_manager: NeoManager) -> None:
        if self.is_focus:
            focus_item = []
        else:
            focus_item = [self.name]
        with neo_manager.driver.session(database=neo_manager.layer.name) as session:
            session.run(
                '''
                MATCH (f:FrameInstance{frame: $name, index: $index})
                SET f.focus = f.focus + $focus_item, f += $name_value_pair
                ''', name=frame.name, index=frame.uuid, focus_item=focus_item, name_value_pair={self.name: self.value}
            )
    
    def apply(self, fun: Callable[[T], T], filter: Callable[[bool, bool], bool] = lambda i, j: True) -> ElementContainer:
        if filter(self.is_core, self.is_focus):
            _container = BaseContainer(self.name, self.is_core, self.type_restraint)
            _container.is_focus = self.is_focus
            _container.value = fun(self.value)
            _container._hash()
            _container._process = Process.Checked
            return _container
        else:
            return None

class SingleContainer(ElementContainer):
    def __init__(self, name, is_core, type_restraint) -> None:
        super().__init__(name, is_core, type_restraint)
        self.value: Union[Frame, Ontology] = None

    def _hash(self) -> ElementContainer:
        _hash = Hash(self.name).update(self.is_focus)
        if self.value.hash is None:
            self.value._hash()
        _hash.update(self.value.hash)
        self.hash = _hash.value()
        return self

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        frame_manager = frame.frame_manager
        if self.name in data['element']:
            _value = data['element'][self.name]
            if isinstance(_value, str):
                self.value = Ontology.create(frame_manager, _value)
            elif isinstance(_value, dict):
                self.value = Frame(frame_manager, _value['frame']).create_from_dict(_value)
            else:
                raise ValueError
            self.is_focus = self.name in data['focus']
            self._process = Process.Filled
            return self
        else:
            return None
    
    def create_from_db(self, frame: Frame, neo_manager: NeoManager, frame_node: Node) -> ElementContainer:
        layer = neo_manager.layer
        with neo_manager.driver.session(database=layer.name) as session:
            res = session.run(
                '''
                MATCH (:FrameInstance{frame: $name, index: $index}) -[r:FrameElement{name: $element_name}]-> (n)
                RETURN r, n
                ''', name = frame_node['frame'], index = frame_node['index'], element_name = self.name
            )
            _res = res.values()
            if _res:
                relationship: Relationship = _res[0][0]
                node: Node = _res[0][1]
                self.is_focus = relationship['is_focus']
                if 'FrameInstance' in node.labels:
                    self.value = Frame(frame.frame_manager, node['frame']).create_from_db(neo_manager, node)
                elif 'Ontology' in node.labels:
                    self.value = Ontology(frame.frame_manager, node['name']).create_from_db(node)
                self._process = Process.Stored
                return self
            else:
                return None

    def save(self, frame: Frame, neo_manager: NeoManager) -> None:
        self.value.save(neo_manager)
        value_uuid = self.value.uuid
        with neo_manager.driver.session(database=neo_manager.layer.name) as session:
            session.run(
                '''
                MATCH (f:FrameInstance{frame: $frame, index: $index}), (e:FrameInstance|Ontology{index: $element_index})
                CREATE (f) -[:FrameElement{name: $element_name, is_focus: $is_focus}]-> (e)
                ''', frame=frame.name, index=frame.uuid, element_index=value_uuid, element_name=self.name, is_focus=self.is_focus
            )

    def apply(self, fun: Callable[[T], T], filter: Callable[[bool, bool], bool] = lambda i, j: True) -> ElementContainer:
        if filter(self.is_core, self.is_focus):
            _container = SingleContainer(self.name, self.is_focus, self.type_restraint)
            _container.is_focus = self.is_focus
            _container.value = fun(self.value)
            _container._hash()
            _container._process = Process.Checked
            return _container
        else:
            return None

class SetContainer(ElementContainer):
    def __init__(self, name, is_core, type_restraint) -> None:
        super().__init__(name, is_core, type_restraint)
        self.value: set[Union[Frame, Ontology]] = set()
        self.focus_items: list[Union[Frame, Ontology]] = []

    def _hash(self) -> ElementContainer:
        _hash = Hash(self.name).update(self.is_focus)
        for _value in self.value:
            if _value.hash is None:
                _value._hash()
        for _value_hash in sorted([_value.hash for _value in self.value]):
            _hash.update(_value_hash)
        for _value_hash in sorted([_value.hash for _value in self.focus_items]):
            _hash.update(_value_hash)
        self.hash = _hash.value()
        return self

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        frame_manager = frame.frame_manager
        if self.name in data['element']:
            _values = data['element'][self.name]
            for _value in _values:
                if isinstance(_value, str):
                    self.value.add(Ontology.create(frame_manager, _value))
                elif isinstance(_value, dict):
                    self.value.add(Frame(frame_manager, _value['frame']).create_from_dict(_value))
                else:
                    raise ValueError
            self.is_focus = self.name in data['containerFocus']
            if self.is_focus:
                for _value in data['containerFocus'][self.name]:
                    if isinstance(_value, str):
                        self.focus_items.append(Ontology.create(frame_manager, _value))
                    elif isinstance(_value, dict):
                        self.focus_items.append(Frame(frame_manager, _value['frame']).create_from_dict(_value))
                    else:
                        raise ValueError
            self._process = Process.Filled
            return self
        else:
            return None

    def create_from_db(self, frame: Frame, neo_manager: NeoManager, frame_node: Node) -> ElementContainer:
        layer = neo_manager.layer
        with neo_manager.driver.session(database=layer.name) as session:
            res = session.run(
                '''
                MATCH (:FrameInstance{frame: $name, index: $index}) -[r:FrameElement{name: $element_name}]-> (n)
                RETURN r, n
                ''', name = frame_node['frame'], index = frame_node['index'], element_name = self.name
            )
            _res = res.values()
            if _res:
                for _element_value in _res:
                    relationship: Relationship = _element_value[0]
                    node: Node = _element_value[1]
                    _is_focus = relationship['is_focus']
                    if 'FrameInstance' in node.labels:
                        _value = Frame(frame.frame_manager, node['frame']).create_from_db(neo_manager, node)
                    elif 'Ontology' in node.labels:
                        _value = Ontology(frame.frame_manager, node['name']).create_from_db(node)
                    self.value.add(_value)
                    if _is_focus:
                        self.is_focus = True
                        self.focus_items.append(_value)
                if self.is_focus is None:
                    self.is_focus = False
                self._process = Process.Stored
                return self
            else:
                return None

    def save(self, frame: Frame, neo_manager: NeoManager) -> None:
        for value in self.value:
            value.save(neo_manager)
        for value in self.value:
            value_uuid = value.uuid
            is_focus = value in self.focus_items
            with neo_manager.driver.session(database=neo_manager.layer.name) as session:
                session.run(
                    '''
                    MATCH (f:FrameInstance{frame: $frame, index: $index}), (e:FrameInstance|Ontology{index: $element_index})
                    CREATE (f) -[:FrameElement{name: $element_name, is_focus: $is_focus}]-> (e)
                    ''', frame=frame.name, index=frame.uuid, element_index=value_uuid, element_name=self.name, is_focus=is_focus
                )

    def apply(self, fun: Callable[[T], T], filter: Callable[[bool, bool], bool] = lambda i, j: True) -> ElementContainer:
        if filter(self.is_core, self.is_focus):
            _container = SetContainer(self.name, self.is_core, self.type_restraint)
            _container.is_focus = self.is_focus
            _container.value = set()
            for _value in self.value:
                _value_is_focus = _value in self.focus_items
                if filter(self.is_core, _value_is_focus):
                    _new_value = fun(_value)
                    _container.value.add(_new_value)
                    if _value_is_focus:
                        _container.focus_items.append(_new_value)
            _container._hash()
            _container._process = Process.Checked
            return _container
        else:
            return None

    def add(self, value: Union[Frame, Ontology], is_focus: bool) -> SetContainer:
        self.value.add(value)
        if is_focus:
            self.is_focus = True
            self.focus_items.append(value)
        return self

class ListContainer(ElementContainer):
    def __init__(self, name, is_core, type_restraint) -> None:
        super().__init__(name, is_core, type_restraint)
        self.value: list[Union[Frame, Ontology]] = []
        self.focus_items: list[Union[Frame, Ontology]] = []

    def _hash(self) -> ElementContainer:
        _hash = Hash(self.name).update(self.is_focus)
        for _value in self.value:
            if _value.hash is None:
                _value._hash()
        for _value in self.value:
            _hash.update(_value.hash)
        for _value_hash in sorted([_value.hash for _value in self.focus_items]):
            _hash.update(_value_hash)
        self.hash = _hash.value()
        return self

    def create_from_dict(self, frame: Frame, data: dict) -> ElementContainer:
        frame_manager = frame.frame_manager
        if self.name in data['element']:
            _values = data['element'][self.name]
            for _value in _values:
                if isinstance(_value, str):
                    self.value.append(Ontology.create(frame_manager, _value))
                elif isinstance(_value, dict):
                    self.value.append(Frame(frame_manager, _value['frame']).create_from_dict(_value))
                else:
                    raise ValueError
            self.is_focus = self.name in data['containerFocus']
            if self.is_focus:
                for _value in data['containerFocus'][self.name]:
                    if isinstance(_value, str):
                        self.focus_items.append(Ontology.create(frame_manager, _value))
                    elif isinstance(_value, dict):
                        self.focus_items.append(Frame(frame_manager, _value['frame']).create_from_dict(_value))
                    else:
                        raise ValueError
            self._process = Process.Filled
            return self
        else:
            return None

    def create_from_db(self, frame: Frame, neo_manager: NeoManager, frame_node: Node) -> ElementContainer:
        layer = neo_manager.layer
        with neo_manager.driver.session(database=layer.name) as session:
            res = session.run(
                '''
                MATCH (:FrameInstance{frame: $name, index: $index}) -[r:FrameElement{name: $element_name}]-> (n)
                RETURN r, n
                ''', name = frame_node['frame'], index = frame_node['index'], element_name = self.name
            )
            _res = res.values()
            if _res:
                relationship: Relationship = _res[0][0]
                start_node: Node = _res[0][1]
                start_value = self.create_structure_from_db(frame.frame_manager, neo_manager, start_node)
                self.value.append(start_value)
                if relationship['is_focus']:
                    self.is_focus = True
                    self.focus_items.append(start_value)
                res = session.run(
                    '''
                    MATCH p=(:FrameInstance|Ontology{index: $first_index})-[:Next*0..]->(n)
                    WHERE all(rel in relationships(p) WHERE rel.index= $frame_index) and not((n)-[:Next{index: $frame_index}]->())
                    with [rel in relationships(p) | [rel, endnode(rel)]] as pairs
                    unwind pairs as pair
                    return pair[0] as relationship, pair[1] as node
                    ''', first_index = start_node['index'], frame_index = frame_node['index']
                )
                _res = res.values()
                for [relationship, node] in _res:
                    value = self.create_structure_from_db(frame.frame_manager, neo_manager, node)
                    self.value.append(value)
                    if relationship['is_focus']:
                        self.is_focus = True
                        self.focus_items.append(value)
                if self.is_focus is None:
                    self.is_focus = False
                self._process = Process.Stored
                return self
            else:
                return None

    def save(self, frame: Frame, neo_manager: NeoManager) -> None:
        for value in self.value:
            value.save(neo_manager)
        first_value_uuid = self.value[0].uuid
        first_is_focus = self.value[0] in self.focus_items
        with neo_manager.driver.session(database=neo_manager.layer.name) as session:
            session.run(
                '''
                MATCH (f:FrameInstance{frame: $frame, index: $index}), (e:FrameInstance|Ontology{index: $element_index})
                CREATE (f) -[:FrameElement{name: $element_name, is_focus: $is_focus}]-> (e)
                ''', frame=frame.name, index=frame.uuid, element_index=first_value_uuid, element_name=self.name, is_focus=first_is_focus
            )
        for value_index, value in enumerate(self.value):
            if value_index != 0:
                start_value_uuid = self.value[value_index-1].uuid
                end_value_uuid = value.uuid
                is_focus = value in self.focus_items
                with neo_manager.driver.session(database=neo_manager.layer.name) as session:
                    session.run(
                        '''
                        MATCH (s:FrameInstance|Ontology{index: $start_index}), (e:FrameInstance|Ontology{index: $end_index})
                        CREATE (s) -[:Next{index: $frame_index, is_focus: $is_focus}]-> (e)
                        ''', start_index=start_value_uuid, end_index=end_value_uuid, frame_index=frame.uuid, is_focus=is_focus
                    )

    def apply(self, fun: Callable[[T], T], filter: Callable[[bool, bool], bool] = lambda i, j: True) -> ElementContainer:
        if filter(self.is_core, self.is_focus):
            _container = ListContainer(self.name, self.is_core, self.type_restraint)
            _container.is_focus = self.is_focus
            _container.value = []
            for _value in self.value:
                _value_is_focus = _value in self.focus_items
                if filter(self.is_core, _value_is_focus):
                    _new_value = fun(_value)
                    _container.value.append(_new_value)
                    if _value_is_focus:
                        _container.focus_items.append(_new_value)
            _container._hash()
            _container._process = Process.Checked
            return _container
        else:
            return None

    def add(self, value: Union[Frame, Ontology], is_focus: bool) -> ListContainer:
        self.value.append(value)
        if is_focus:
            self.is_focus = True
            self.focus_items.append(value)
        return self

class Ontology:
    def __init__(self, frame_manager: FrameManager, name: str) -> None:
        self.name = name
        self.marker = name
        self.frame_manager = frame_manager
        self.frame_manager.marker_map[name] = self
        self.hash = Hash(name).value()
        self.uuid = uuid4().hex
        self.frame_manager.marker_map[self.uuid] = self
        self.frame_manager.marker_map_primary[self.uuid] = self
        self.frame_manager.marker_map_prototype[self.uuid] = self
        self.corpus_ids = []

    def create_from_db(self, node: Node):
        self.uuid = node['index']
        return self

    def save(self, neo_manager: NeoManager):
        with neo_manager.driver.session(database=neo_manager.layer.name) as session:
            res = session.run(
                '''
                MATCH (o:Ontology{name: $name})
                RETURN o.index
                ''', name=self.name
            )
            _res = res.value()
            if _res:
                self.uuid = _res[0]
            else:
                session.run(
                    '''
                    CREATE (:Ontology{name: $name, index: $index, hash: $hash})
                    ''', name=self.name, index=self.uuid, hash=self.hash
                )


    @staticmethod
    def create(frame_manager: FrameManager, name: str) -> Ontology:
        if name in frame_manager.marker_map:
            return frame_manager.marker_map[name]
        else:
            return Ontology(frame_manager, name)

    def append_corpus_id(self, corpus_id: str) -> Ontology:
        if corpus_id not in self.corpus_ids:
            self.corpus_ids.append(corpus_id)
        return self